{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import scipy\n",
    "import xarray as xr\n",
    "from skdownscale.pointwise_models import AnalogRegression, PureAnalog"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# open a small dataset for training\n",
    "training = xr.open_zarr(\"../data/downscale_test_data.zarr.zip\", group=\"training\")\n",
    "training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# open a small dataset of observations (targets)\n",
    "targets = xr.open_zarr(\"../data/downscale_test_data.zarr.zip\", group=\"targets\")\n",
    "targets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# extract 1 point of training data for precipitation and temperature\n",
    "X_temp = training.isel(point=0).to_dataframe()[[\"T2max\"]] - 273.13\n",
    "X_pcp = training.isel(point=0).to_dataframe()[[\"PREC_TOT\"]] * 24\n",
    "display(X_temp.head(), X_pcp.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# extract 1 point of target data for precipitation and temperature\n",
    "y_temp = targets.isel(point=0).to_dataframe()[[\"Tmax\"]]\n",
    "y_pcp = targets.isel(point=0).to_dataframe()[[\"Prec\"]]\n",
    "display(y_temp.head(), y_pcp.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fit/predict using the PureAnalog class\n",
    "for kind in [\"best_analog\", \"sample_analogs\", \"weight_analogs\", \"mean_analogs\"]:\n",
    "    pure_analog = PureAnalog(kind=kind, n_analogs=10)\n",
    "    pure_analog.fit(X_temp[:1000], y_temp[:1000])\n",
    "    out = pure_analog.predict(X_temp[1000:])\n",
    "\n",
    "    plt.plot(out[:300], label=kind)\n",
    "\n",
    "# Fit/predict using the AnalogRegression class\n",
    "analog_reg = AnalogRegression(n_analogs=100)\n",
    "analog_reg.fit(X_temp[:1000], y_temp[:1000])\n",
    "out = analog_reg.predict(X_temp[1000:])\n",
    "plt.plot(out[:300], label=\"AnalogRegression\")\n",
    "plt.legend()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
